package com.starhub.application.rag.util;

import com.starhub.application.rag.constant.RagConstant;
import com.starhub.utils.Util;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;

/**
 * 搜索结果工具类
 * @version: V1.0
 * @author: liuhf
 * @time 2025-06-05
 */
@Slf4j
@Component
public class EmbeddingSearchResultUtil {

    /**
     * 将搜索结果转换为TextSegment列表
     * @param searchResult
     * @param key
     * @return
     */
    public static List<TextSegment> transToSegments(EmbeddingSearchResult<TextSegment> searchResult) {
        return searchResult.matches().stream()
                .map(match -> match.embedded())
                .collect(Collectors.toList());
    }

    /**
     * 将搜索结果转换为段落ID列表
     * @param searchResult
     * @param key
     * @return
     */
    public static List<String> transToSegmentIds(EmbeddingSearchResult<TextSegment> searchResult) {
        return transToIds(searchResult, RagConstant.EmbeddingMetaData.SEGMENT_ID);
    }
    
    
    /**
     * 将搜索结果转换为知识条目ID列表
     *
     * @param searchResult 
     * @param key 来源RagConstant.EmbeddingMetaData 
     * @return 知识条目ID列表
     */
    public static List<String> transToIds(EmbeddingSearchResult<TextSegment> searchResult, String key) {
        return searchResult.matches().stream()
                .map(match -> convertMatchToItemId(match, key))
                .collect(Collectors.toList());

    }

    /**
     * 从匹配结果中提取知识条目ID
     *
     * @param match 匹配结果
     * @param key 来源RagConstant.EmbeddingMetaData 
     * @return 知识条目ID
     */
    private static String convertMatchToItemId(EmbeddingMatch<TextSegment> match, String key) {
        if (match == null || match.embedded() == null) {
            return null;
        }

        try {
            Metadata metadata = match.embedded().metadata();
            return metadata != null ? Util.null2String(metadata.get(key)) : null;
        } catch (Exception e) {
            log.error("提取知识条目ID失败: {}", e.getMessage(), e);
            return null;
        }
    }
}